package com.matrix.framework.core.component;

import com.matrix.framework.auth.data.UserPo;
import com.matrix.framework.auth.service.UserService;
import com.matrix.framework.core.common.global.SpringUtil;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureAlgorithm;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;

import javax.crypto.SecretKey;
import java.util.Date;
import java.util.Map;

/**
 * JWT工具类 - 支持密钥轮换机制
 *
 * Copyright © 海平面工作室 版权所有
 *
 * @Author: Leo
 * @Create: 2024/10/3 11:15
 * @Since 1.2
 */
@Component
public class Jwt {

    private final JwtKeyManager keyManager;
    
    // JWT过期时间（30分钟）- 默认值，可通过配置动态修改
    public static long JWT_EXPIRATION = 1000 * 60 * 30;
    
    // 自动续期阈值时间（10分钟，即当剩余时间少于15分钟时自动续期）- 默认值，可通过配置动态修改
    public static long REFRESH_THRESHOLD = 1000 * 60 * 10;

    public Jwt(JwtKeyManager keyManager) {
        this.keyManager = keyManager;
    }

    /**
     * 设置JWT过期时间
     * @param expirationMinutes 过期时间（分钟）
     */
    public static void setJwtExpiration(long expirationMinutes) {
        JWT_EXPIRATION = expirationMinutes * 60 * 1000;
    }

    /**
     * 设置JWT刷新阈值时间
     * @param thresholdMinutes 阈值时间（分钟）
     */
    public static void setRefreshThreshold(long thresholdMinutes) {
        REFRESH_THRESHOLD = thresholdMinutes * 60 * 1000;
    }

    // 生成JWT - 包含keyId
    public String generateToken(String id, String username) {
        String currentKeyId = keyManager.getCurrentKeyId();
        SecretKey currentKey = keyManager.getCurrentKey();
        
        return Jwts.builder()
                .setId(id)
                .setSubject(username)  // 设置用户名
                .setIssuedAt(new Date())  // 设置签发时间
                .setExpiration(new Date(System.currentTimeMillis() + JWT_EXPIRATION))  // 设置过期时间
                .claim("keyId", currentKeyId)  // 添加密钥ID
                .signWith(currentKey, SignatureAlgorithm.HS256)  // 使用当前密钥签名
                .compact();  // 生成JWT
    }

    // 验证Token - 支持多密钥验证
    public Boolean validateToken(String token, String username) {
        try {
            final String extractedUsername = extractUsername(token);
            return (extractedUsername.equals(username) && !isTokenExpired(token));
        } catch (Exception e) {
            return false;
        }
    }
    
    // 验证Token并返回验证结果信息 - 支持多密钥验证
    public TokenValidationResult validateTokenWithResult(String token, String username) {
        try {
            final String extractedUsername = extractUsername(token);
            if (!extractedUsername.equals(username)) {
                return new TokenValidationResult(false, false, null);
            }
            
            boolean isExpired = isTokenExpired(token);
            if (isExpired) {
                return new TokenValidationResult(false, false, null);
            }
            
            // 检查是否需要续期
            boolean needsRefresh = needsRefresh(token);
            String newToken = null;
            if (needsRefresh) {
                String userId = extractUserId(token);
                newToken = generateToken(userId, username);
            }
            
            return new TokenValidationResult(true, needsRefresh, newToken);
        } catch (Exception e) {
            return new TokenValidationResult(false, false, null);
        }
    }

    // 提取用户名 - 支持多密钥验证
    public String extractUsername(String token) {
        Map<String, SecretKey> validKeys = keyManager.getValidKeys();
        
        for (Map.Entry<String, SecretKey> entry : validKeys.entrySet()) {
            try {
                return Jwts.parserBuilder()
                        .setSigningKey(entry.getValue())
                        .build()
                        .parseClaimsJws(token)
                        .getBody()
                        .getSubject();
            } catch (Exception e) {
                // 继续尝试下一个密钥
                continue;
            }
        }
        
        throw new RuntimeException("无法验证JWT token");
    }

    // 提取用户ID - 支持多密钥验证
    public String extractUserId(String token) {
        Map<String, SecretKey> validKeys = keyManager.getValidKeys();
        
        for (Map.Entry<String, SecretKey> entry : validKeys.entrySet()) {
            try {
                return Jwts.parserBuilder()
                        .setSigningKey(entry.getValue())
                        .build()
                        .parseClaimsJws(token)
                        .getBody()
                        .getId();
            } catch (Exception e) {
                // 继续尝试下一个密钥
                continue;
            }
        }
        
        throw new RuntimeException("无法验证JWT token");
    }

    // 检查Token是否过期 - 支持多密钥验证
    private Boolean isTokenExpired(String token) {
        try {
            return extractExpiration(token).before(new Date());
        } catch (Exception e) {
            return true;
        }
    }

    // 提取过期时间 - 支持多密钥验证
    public Date extractExpiration(String token) {
        Map<String, SecretKey> validKeys = keyManager.getValidKeys();
        
        for (Map.Entry<String, SecretKey> entry : validKeys.entrySet()) {
            try {
                return Jwts.parserBuilder()
                        .setSigningKey(entry.getValue())
                        .build()
                        .parseClaimsJws(token)
                        .getBody()
                        .getExpiration();
            } catch (Exception e) {
                // 继续尝试下一个密钥
                continue;
            }
        }
        
        throw new RuntimeException("无法验证JWT token");
    }
    
    // 检查Token是否需要续期（剩余时间少于阈值时返回true）
    public Boolean needsRefresh(String token) {
        try {
            Date expiration = extractExpiration(token);
            long remainingTime = expiration.getTime() - System.currentTimeMillis();
            return remainingTime <= REFRESH_THRESHOLD;
        } catch (Exception e) {
            return true;
        }
    }
    
    // 获取Token的剩余有效时间（毫秒）
    public long getRemainingTime(String token) {
        try {
            Date expiration = extractExpiration(token);
            return expiration.getTime() - System.currentTimeMillis();
        } catch (Exception e) {
            return 0;
        }
    }

    // 获取登录用户名
    public static String getLoginUserName(ServerWebExchange exchange) {
        String token = exchange.getRequest().getHeaders().getFirst("Authorization");
        if (token == null) {
            return null;
        }
        Jwt jwt = SpringUtil.getBean(Jwt.class);
        return jwt.extractUsername(token);
    }

    // 获取登录用户ID
    public static long getLoginUserId(ServerWebExchange exchange) {
        String token = exchange.getRequest().getHeaders().getFirst("Authorization");
        if (token == null) {
            return 0;
        }
        Jwt jwt = SpringUtil.getBean(Jwt.class);
        return Long.parseLong(jwt.extractUserId(token));
    }

    // 获取登录用户
    public static UserPo getLoginUserByName(ServerWebExchange exchange) {
        String token = exchange.getRequest().getHeaders().getFirst("Authorization");
        if (token == null) {
            return null;
        }
        Jwt jwt = SpringUtil.getBean(Jwt.class);
        String userName = jwt.extractUsername(token);
        UserService userService = SpringUtil.getBean(UserService.class);
        return userService.getUserByName(userName).block();
    }

    // 获取登录用户
    public static UserPo getLoginUserById(ServerWebExchange exchange) {
        String token = exchange.getRequest().getHeaders().getFirst("Authorization");
        if (token == null) {
            return null;
        }
        Jwt jwt = SpringUtil.getBean(Jwt.class);
        String userId = jwt.extractUserId(token);
        UserService userService = SpringUtil.getBean(UserService.class);
        return userService.getUserById(Long.parseLong(userId)).block();
    }
    
    /**
     * Token验证结果类
     */
    public static class TokenValidationResult {
        private final boolean valid;          // Token是否有效
        private final boolean needsRefresh;   // 是否需要续期
        private final String newToken;        // 新的Token（如果需要续期）
        
        public TokenValidationResult(boolean valid, boolean needsRefresh, String newToken) {
            this.valid = valid;
            this.needsRefresh = needsRefresh;
            this.newToken = newToken;
        }
        
        public boolean isValid() {
            return valid;
        }
        
        public boolean needsRefresh() {
            return needsRefresh;
        }
        
        public String getNewToken() {
            return newToken;
        }
    }
}

